import numpy as np
import scipy.signal as signal
import matplotlib.pyplot as plt

plt.rcParams['font.sans-serif']=['SimHei']

def transfer_function(J, q, ts):
    sys = signal.TransferFunction(1, [J, q, 0])
    dsys = signal.cont2discrete((1, [J, q, 0]), ts, method='zoh')
    return dsys

def calculate(dsys, len_time, ts):
    num = dsys[0][0]
    den = dsys[1]
    u_1 = 0
    u_2 = 0
    y_1 = 0
    y_2 = 0
    error_1 = 0
    ei = 0
    yd = 0.5*np.ones(len_time)
    y = np.zeros(len_time)
    error = np.zeros(len_time)
    derror = np.zeros(len_time)
    kp = np.zeros(len_time)
    kd = np.zeros(len_time)
    ki = np.zeros(len_time)
    u = np.zeros(len_time)

    for k in range(len_time):
        y[k] = -den[1] * y_1 - den[2] * y_2 + num[1] * u_1 + num[2] * u_2
        error[k] = yd[k] - y[k]
        derror[k] = (error[k] - error_1) / ts

        # Define your constants here
        ap = 22
        bp = 8.0
        cp = 0.8
        kp[k] = ap + bp * (1 - 1 / np.cosh(cp * error[k]))

        ad = 0.5
        bd = 2.5
        cd = 6.5
        dd = 0.30
        kd[k] = ad + bd / (1 + cd * np.exp(dd * error[k]))

        ai = 1
        ci = 1
        ki[k] = ai * np.cosh(ci * error[k])

        ei = ei + error[k] * ts
        u[k] = kp[k] * error[k] + kd[k] * derror[k] + ki[k] * ei

        u_2 = u_1
        u_1 = u[k]
        y_2 = y_1
        y_1 = y[k]
        error_1 = error[k]

    return yd, y, error, kp, kd, ki

def plot_results(time, yd, y, error, kp, kd, ki):
    plt.figure(1)
    plt.plot(time, yd, 'r', time, y, 'k', linewidth=2)
    plt.xlabel('时间(s)')
    plt.ylabel('位置信号')
    plt.legend(['上顶栓压力设定值', '实际响应'], loc='upper right')

    plt.figure(2)
    plt.subplot(311)
    plt.plot(error, kp, 'r', linewidth=2)
    plt.xlabel('error')
    plt.ylabel('kp')
    plt.subplot(312)
    plt.plot(error, kd, 'r', linewidth=2)
    plt.xlabel('error')
    plt.ylabel('kd')
    plt.subplot(313)
    plt.plot(error, ki, 'r', linewidth=2)
    plt.xlabel('error')
    plt.ylabel('ki')

    plt.figure(3)
    plt.subplot(311)
    plt.plot(time, kp, 'r', linewidth=2)
    plt.xlabel('time(s)')
    plt.ylabel('kp')
    plt.subplot(312)
    plt.plot(time, kd, 'r', linewidth=2)
    plt.xlabel('time(s)')
    plt.ylabel('kd')
    plt.subplot(313)
    plt.plot(time, ki, 'r', linewidth=2)
    plt.xlabel('time(s)')
    plt.ylabel('ki')

    plt.show()

if __name__ == "__main__":
    ts = 0.004
    J = 2 / 133
    q = 25 / 133
    dsys = transfer_function(J, q, ts)

    len_time = 500
    time = np.arange(0, len_time * ts, ts)
    yd, y, error, kp, kd, ki = calculate(dsys, len_time, ts)

    plot_results(time, yd, y, error, kp, kd, ki)